#!/usr/bin/env python3
"""
医学图像评估工具：计算SSIM和DICE指标
支持多种对比方法的评估，包括：
1. 图像重采样 (reslice)
2. 计算图像质量指标 (SSIM)
3. 计算分割精度指标 (DICE)
"""

import os
import SimpleITK as sitk
import numpy as np
from skimage.metrics import structural_similarity as ssim
from glob import glob
from scipy import ndimage as nd
import csv
import argparse
from pathlib import Path
from typing import Dict, List, Optional, Tuple
import logging

# 配置日志
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)


class MedicalImageProcessor:
    """医学图像处理器"""

    def __init__(self, base_dir: str):
        self.base_dir = Path(base_dir)
        self.logger = logging.getLogger(__name__)

        # 定义目录结构
        self.dirs = {
            "prediction": self.base_dir / "prediction",
            "reference": self.base_dir / "reference",
            "segmentation": self.base_dir / "segmentation",
            "resliced": self.base_dir / "resliced",
            "metric_results": self.base_dir / "metric_results",
        }

        # 确保输出目录存在
        self.dirs["metric_results"].mkdir(exist_ok=True)

    def get_subjects(self) -> List[str]:
        """获取所有受试者ID"""
        if not self.dirs["reference"].exists():
            return []
        return [d.name for d in self.dirs["reference"].iterdir() if d.is_dir()]

    def _get_reference_paths(
        self, subject_id: str
    ) -> Tuple[Optional[Path], Optional[Path]]:
        """
        根据受试者ID动态确定参考图像和掩膜路径

        Args:
            subject_id: 受试者ID

        Returns:
            (ref_image_path, ref_mask_path) 元组，如果文件不存在则返回None
        """
        ref_dir = self.dirs["reference"] / subject_id

        # 根据subject_id后缀确定方向
        if "_tra" in subject_id:
            # 轴向 (transverse/axial)
            orientation = "axial"
        elif "_sag" in subject_id:
            # 矢状面 (sagittal)
            orientation = "sagittal"
        else:
            # 默认尝试轴向，如果不存在则尝试矢状面
            orientation = "axial"

        # 构建文件路径
        ref_image_path = ref_dir / f"source_{orientation}_LR_registered.nii.gz"
        ref_mask_path = ref_dir / f"prostate_source_{orientation}_LR_registered.nii.gz"

        # 如果默认方向的文件不存在，尝试另一个方向
        if not ref_image_path.exists() or not ref_mask_path.exists():
            alternative_orientation = "sagittal" if orientation == "axial" else "axial"
            alt_image_path = (
                ref_dir / f"source_{alternative_orientation}_LR_registered.nii.gz"
            )
            alt_mask_path = (
                ref_dir
                / f"prostate_source_{alternative_orientation}_LR_registered.nii.gz"
            )

            if alt_image_path.exists() and alt_mask_path.exists():
                self.logger.info(
                    f"使用 {alternative_orientation} 方向的参考文件 (主方向 {orientation} 不可用)"
                )
                return alt_image_path, alt_mask_path

        # 检查文件是否存在
        if ref_image_path.exists() and ref_mask_path.exists():
            self.logger.info(f"使用 {orientation} 方向的参考文件")
            return ref_image_path, ref_mask_path
        else:
            # 列出可用的参考文件以便调试
            available_files = list(ref_dir.glob("*.nii.gz")) if ref_dir.exists() else []
            self.logger.error(f"未找到受试者 {subject_id} 的参考文件")
            self.logger.error(
                f"期望的文件: {ref_image_path.name}, {ref_mask_path.name}"
            )
            if available_files:
                self.logger.error(f"可用的文件: {[f.name for f in available_files]}")
            return None, None

    def create_subject_directories(
        self, subject_ids: List[str], dry_run: bool = False
    ) -> bool:
        """
        为指定的受试者自动创建标准文件夹结构

        目录结构:
        ├── prediction/
        │   ├── subject1/
        │   │   ├── processed/
        │   │   ├── proposed/
        │   │   ├── source/
        │   │   └── contrast_methods/
        │   └── subject2/...
        ├── segmentation/
        │   ├── subject1/
        │   └── subject2/...
        ├── reference/
        │   ├── subject1/
        │   └── subject2/
        ├── resliced/
        │   ├── subject1/
        │   │   ├── images/
        │   │   └── masks/

        Args:
            subject_ids: 受试者ID列表
            dry_run: 是否为模拟运行，True时只显示将要创建的目录

        Returns:
            创建是否成功
        """
        if not subject_ids:
            self.logger.warning("未提供受试者ID列表")
            return False

        # 定义每个主目录下需要创建的子目录结构
        directory_structure = {
            "prediction": ["processed", "proposed", "source", "contrast_methods"],
            "segmentation": [],  # 只创建subject目录
            "reference": [],  # 只创建subject目录
            "resliced": ["images", "masks"],
        }

        self.logger.info(f"{'='*20} 创建受试者目录结构 {'='*20}")
        self.logger.info(f"受试者列表: {', '.join(subject_ids)}")

        if dry_run:
            self.logger.info("[模拟运行] 将要创建以下目录结构:")

        success = True
        created_dirs = []

        for subject_id in subject_ids:
            self.logger.info(f"--- 处理受试者: {subject_id} ---")

            for main_dir, subdirs in directory_structure.items():
                # 创建主受试者目录
                subject_dir = self.dirs[main_dir] / subject_id

                if dry_run:
                    self.logger.info(f"  [模拟] 创建目录: {subject_dir}")
                else:
                    try:
                        subject_dir.mkdir(parents=True, exist_ok=True)
                        created_dirs.append(str(subject_dir))
                        self.logger.info(f"  ✓ 创建目录: {subject_dir}")
                    except Exception as e:
                        self.logger.error(f"  ✗ 创建目录失败 {subject_dir}: {e}")
                        success = False
                        continue

                # 创建子目录
                for subdir in subdirs:
                    sub_path = subject_dir / subdir

                    if dry_run:
                        self.logger.info(f"    [模拟] 创建子目录: {sub_path}")
                    else:
                        try:
                            sub_path.mkdir(parents=True, exist_ok=True)
                            created_dirs.append(str(sub_path))
                            self.logger.info(f"    ✓ 创建子目录: {sub_path}")
                        except Exception as e:
                            self.logger.error(f"    ✗ 创建子目录失败 {sub_path}: {e}")
                            success = False

        if dry_run:
            self.logger.info(
                f"[模拟运行] 总共将创建 {len([s for s in subject_ids for _ in range(len(directory_structure) + sum(len(subdirs) for subdirs in directory_structure.values()))])} 个目录"
            )
        else:
            self.logger.info(f"目录创建完成! 总共创建了 {len(created_dirs)} 个目录")
            if not success:
                self.logger.warning("部分目录创建失败，请检查权限和路径")

        return success

    def get_method_types(self, subject_id: str) -> Dict[str, List[str]]:
        """获取指定受试者的所有方法类型"""
        methods = {"contrast_methods": [], "proposed": [], "processed": []}

        pred_dir = self.dirs["prediction"] / subject_id
        if pred_dir.exists():
            for method_type in methods.keys():
                method_dir = pred_dir / method_type
                if method_dir.exists():
                    methods[method_type] = [f.name for f in method_dir.glob("*.nii.gz")]

        return methods

    def reslice_image(
        self, image_path: str, reference_image_path: str
    ) -> Optional[sitk.Image]:
        """
        使用参考图像的空间信息对给定图像进行重采样

        Args:
            image_path: 需要重采样的图像路径
            reference_image_path: 参考图像路径

        Returns:
            重采样后的图像对象，失败时返回None
        """
        try:
            image = sitk.ReadImage(str(image_path))
            reference_image = sitk.ReadImage(str(reference_image_path))
        except Exception as e:
            self.logger.error(f"读取图像文件出错: {e}")
            return None

        reslice_filter = sitk.ResampleImageFilter()
        reslice_filter.SetReferenceImage(reference_image)
        reslice_filter.SetInterpolator(sitk.sitkLinear)
        resliced_image = reslice_filter.Execute(image)
        return resliced_image

    def reslice_predictions(
        self,
        subject_id: str,
        method_types: List[str] = None,
        ref_image_path: str = None,
        dry_run: bool = False,
    ) -> bool:
        """
        重采样预测图像

        Args:
            subject_id: 受试者ID
            method_types: 要处理的方法类型列表，None表示处理所有类型
            ref_image_path: reslice参考图像
            dry_run: 是否为模拟运行

        Returns:
            处理是否成功
        """
        # 参考图像路径已经存在
        if not ref_image_path.exists():
            self.logger.error(f"未找到受试者 {subject_id} 的参考图像: {ref_image_path}")
            return False

        # 创建输出目录
        output_dir = self.dirs["resliced"] / subject_id / "images"
        output_dir.mkdir(parents=True, exist_ok=True)

        # 获取所有方法类型
        available_methods = self.get_method_types(subject_id)
        if method_types is None:
            method_types = list(available_methods.keys())

        success = True
        for method_type in method_types:
            method_dir = self.dirs["prediction"] / subject_id / method_type
            if not method_dir.exists():
                self.logger.warning(f"方法目录不存在: {method_dir}")
                continue

            # 处理该方法下的所有文件，直接保存到images目录
            files = list(method_dir.glob("*.nii.gz"))
            if not files:
                self.logger.warning(f"在 {method_dir} 中未找到.nii.gz文件")
                continue

            self.logger.info(f"处理 {method_type}: 找到 {len(files)} 个文件")

            for file_path in files:
                # 直接保存到images目录，不创建方法类型子目录
                output_path = output_dir / file_path.name

                if dry_run:
                    self.logger.info(f"[模拟] 将重采样: {file_path} -> {output_path}")
                else:
                    resliced_img = self.reslice_image(
                        str(file_path), str(ref_image_path)
                    )
                    if resliced_img:
                        sitk.WriteImage(resliced_img, str(output_path))
                        self.logger.info(f"已重采样并保存: {file_path.name}")
                    else:
                        self.logger.error(f"重采样失败: {file_path.name}")
                        success = False

        return success

    def reslice_segmentations(
        self, subject_id: str, ref_mask_path: str, dry_run: bool = False
    ) -> bool:
        """
        重采样分割掩膜

        Args:
            subject_id: 受试者ID
            ref_mask_path: reslice参考掩膜
            dry_run: 是否为模拟运行

        Returns:
            处理是否成功
        """
        # 获取参考掩膜路径
        if not ref_mask_path.exists():
            self.logger.error(f"未找到受试者 {subject_id} 的参考掩膜: {ref_mask_path}")
            return False

        # 分割掩膜目录
        seg_dir = self.dirs["segmentation"] / subject_id
        if not seg_dir.exists():
            self.logger.warning(f"分割目录不存在: {seg_dir}")
            return False

        # 创建输出目录
        output_dir = self.dirs["resliced"] / subject_id / "masks"
        output_dir.mkdir(parents=True, exist_ok=True)

        # 处理根目录下的.nii.gz文件
        files = list(seg_dir.glob("*.nii.gz"))
        if not files:
            self.logger.warning(f"在 {seg_dir} 中未找到.nii.gz文件")
            return False

        self.logger.info(f"处理分割掩膜: 找到 {len(files)} 个文件")

        success = True
        for file_path in files:
            output_path = output_dir / file_path.name

            if dry_run:
                self.logger.info(f"[模拟] 将重采样掩膜: {file_path} -> {output_path}")
            else:
                resliced_img = self.reslice_image(str(file_path), str(ref_mask_path))
                if resliced_img:
                    sitk.WriteImage(resliced_img, str(output_path))
                    self.logger.info(f"已重采样并保存掩膜: {file_path.name}")
                else:
                    self.logger.error(f"重采样掩膜失败: {file_path.name}")
                    success = False

        return success

    def calculate_ssim(self, image1_path: str, image2_path: str) -> Optional[float]:
        """
        计算两幅图像之间的结构相似性指数 (SSIM)

        Args:
            image1_path: 第一幅图像路径，参考图像
            image2_path: 第二幅图像路径，待计算图像

        Returns:
            SSIM值，保留四位小数。如果失败则返回None
        """
        try:
            img1 = sitk.ReadImage(image1_path)
            img2 = sitk.ReadImage(image2_path)
        except Exception as e:
            self.logger.error(f"读取SSIM图像文件出错: {e}")
            return None

        img1_array = sitk.GetArrayFromImage(img1).astype(np.float32)
        img2_array = sitk.GetArrayFromImage(img2).astype(np.float32)

        if img1_array.shape != img2_array.shape:
            self.logger.warning(
                f"图像尺寸不一致 {img1_array.shape} vs {img2_array.shape}，无法计算SSIM"
            )
            return None

        # 创建掩膜，过滤背景
        mask = img2_array > 0

        # 对掩膜应用形态学操作以改进质量
        struct = nd.generate_binary_structure(3, 1)  # 3D结构
        mask = nd.binary_erosion(mask, structure=struct).astype(mask.dtype)
        mask = nd.binary_fill_holes(mask).astype(mask.dtype)
        mask = nd.binary_dilation(mask, structure=struct).astype(mask.dtype)

        img1_masked = img1_array * mask
        img2_masked = img2_array * mask

        return np.round(ssim(img1_masked, img2_masked, data_range=255), 4)

    def calculate_dice(self, pred_path: str, label_path: str) -> Optional[float]:
        """
        计算预测掩膜和真实掩膜之间的Dice系数

        Args:
            pred_path: 预测掩膜路径
            label_path: 真实掩膜路径

        Returns:
            Dice系数，保留四位小数。如果失败则返回None
        """
        try:
            pred = sitk.ReadImage(pred_path)
            label = sitk.ReadImage(label_path)

            # 确保空间一致性
            if pred.GetSize() != label.GetSize():
                self.logger.warning(
                    f"掩膜尺寸不一致 {pred.GetSize()} vs {label.GetSize()}"
                )
                return None

            pred_array = sitk.GetArrayFromImage(pred).astype(bool)
            label_array = sitk.GetArrayFromImage(label).astype(bool)

        except Exception as e:
            self.logger.error(f"读取Dice掩膜文件出错: {e}")
            return None

        # 如果两个掩膜都是空的，则返回1.0
        if not np.any(label_array) and not np.any(pred_array):
            return 1.0

        # 计算交集和并集
        intersection = np.logical_and(pred_array, label_array).sum()
        total = pred_array.sum() + label_array.sum()

        # 避免除零错误
        if total == 0:
            return 0.0

        # 计算Dice系数
        dice = 2 * intersection / total
        return round(dice, 4)

    def compute_ssim_metrics(self, subject_id: str, ref_image_path: Path) -> bool:
        """
        计算SSIM指标并保存到CSV文件

        Args:
            subject_id: 受试者ID
            ref_image_path: 参考图像路径

        Returns:
            计算是否成功
        """
        if not ref_image_path.exists():
            self.logger.error(f"未找到参考图像: {ref_image_path}")
            return False

        # 重采样图像目录
        resliced_images_dir = self.dirs["resliced"] / subject_id / "images"
        if not resliced_images_dir.exists():
            self.logger.error(f"未找到重采样图像目录: {resliced_images_dir}")
            return False

        # 准备CSV数据
        csv_data = []
        output_csv_path = self.dirs["metric_results"] / f"{subject_id}_ssim_results.csv"

        self.logger.info(f"开始计算 {subject_id} 的SSIM指标")

        # 处理images目录下的所有文件
        files = sorted(resliced_images_dir.glob("*.nii.gz"))
        if not files:
            self.logger.warning(f"在 {resliced_images_dir} 中未找到.nii.gz文件")
            return False

        self.logger.info(f"找到 {len(files)} 个重采样图像文件")

        for file_path in files:
            # 从文件名提取方法名
            method_name = self._extract_method_name(file_path.name)

            ssim_value = self.calculate_ssim(str(ref_image_path), str(file_path))

            if ssim_value is not None:
                csv_data.append(
                    {
                        "method": method_name,
                        "file": file_path.name,
                        "SSIM": f"{ssim_value:.4f}",
                    }
                )
                self.logger.info(
                    f"  {file_path.name} ({method_name}): {ssim_value:.4f}"
                )

        # 写入CSV文件
        if csv_data:
            with open(output_csv_path, "w", newline="", encoding="utf-8") as csvfile:
                fieldnames = ["method", "file", "SSIM"]
                writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
                writer.writeheader()
                writer.writerows(csv_data)
            self.logger.info(f"SSIM结果已保存到 {output_csv_path}")
            return True
        else:
            self.logger.warning("没有SSIM数据可写入")
            return False

    def compute_dice_metrics(self, subject_id: str, ref_mask_path: Path) -> bool:
        """
        计算Dice指标并保存到CSV文件

        Args:
            subject_id: 受试者ID
            ref_mask_path: 参考掩膜路径

        Returns:
            计算是否成功
        """
        if not ref_mask_path.exists():
            self.logger.error(f"未找到参考掩膜: {ref_mask_path}")
            return False

        # 重采样掩膜目录
        resliced_masks_dir = self.dirs["resliced"] / subject_id / "masks"
        if not resliced_masks_dir.exists():
            self.logger.error(f"未找到重采样掩膜目录: {resliced_masks_dir}")
            return False

        # 准备CSV数据
        csv_data = []
        output_csv_path = self.dirs["metric_results"] / f"{subject_id}_dice_results.csv"

        self.logger.info(f"开始计算 {subject_id} 的Dice指标")

        files = sorted(resliced_masks_dir.glob("*.nii.gz"))
        if not files:
            self.logger.warning(f"在 {resliced_masks_dir} 中未找到.nii.gz文件")
            return False

        for file_path in files:
            dice_value = self.calculate_dice(str(file_path), str(ref_mask_path))

            if dice_value is not None:
                # 从文件名推断方法名
                method_name = self._extract_method_name(file_path.name)
                csv_data.append(
                    {
                        "method": method_name,
                        "file": file_path.name,
                        "Dice": f"{dice_value:.4f}",
                    }
                )
                self.logger.info(f"  {file_path.name}: {dice_value:.4f}")

        # 写入CSV文件
        if csv_data:
            with open(output_csv_path, "w", newline="", encoding="utf-8") as csvfile:
                fieldnames = ["method", "file", "Dice"]
                writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
                writer.writeheader()
                writer.writerows(csv_data)
            self.logger.info(f"Dice结果已保存到 {output_csv_path}")
            return True
        else:
            self.logger.warning("没有Dice数据可写入")
            return False

    def _extract_method_name(self, filename: str) -> str:
        """从文件名中提取方法名"""
        # 根据文件命名规则提取方法名
        if "ResUNet2D_FFC_wongf" in filename:
            return "ResUNet2D_FFC_wongf"
        elif "ResUNet2D_FFC" in filename:
            return "ResUNet2D_FFC"
        elif "ResUNet2D" in filename:
            return "ResUNet2D"
        elif "UNet2D" in filename:
            return "UNet2D"
        elif "mcinr" in filename:
            return "mcinr"
        elif "smore" in filename:
            return "smore"
        elif "sagittal" in filename:
            return "sagittal"
        elif "axial" in filename:
            return "axial"
        else:
            return "unknown"

    def process_subject(
        self,
        subject_id: str,
        reslice: bool = True,
        calc_ssim: bool = True,
        calc_dice: bool = True,
        method_types: List[str] = None,
        dry_run: bool = False,
    ) -> bool:
        """
        处理单个受试者的数据

        Args:
            subject_id: 受试者ID
            reslice: 是否执行重采样
            calc_ssim: 是否计算SSIM
            calc_dice: 是否计算Dice
            method_types: 要处理的方法类型列表
            dry_run: 是否为模拟运行

        Returns:
            处理是否成功
        """
        self.logger.info(f"{'='*20} 处理受试者: {subject_id} {'='*20}")

        # 动态获取参考文件路径
        ref_image_path, ref_mask_path = self._get_reference_paths(subject_id)

        if ref_image_path is None or ref_mask_path is None:
            return False

        success = True

        # 1. 执行重采样
        if reslice:
            self.logger.info("--- 步骤1: 图像重采样 ---")
            if not self.reslice_predictions(
                subject_id, method_types, ref_image_path, dry_run
            ):
                success = False

            if not self.reslice_segmentations(subject_id, ref_mask_path, dry_run):
                success = False

        # 2. 计算SSIM
        if calc_ssim and not dry_run:
            self.logger.info("--- 步骤2: 计算SSIM ---")
            if not self.compute_ssim_metrics(subject_id, ref_image_path=ref_image_path):
                success = False

        # 3. 计算Dice
        if calc_dice and not dry_run:
            self.logger.info("--- 步骤3: 计算Dice ---")
            if not self.compute_dice_metrics(subject_id, ref_mask_path=ref_mask_path):
                success = False

        return success


def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(
        description="医学图像评估工具: 计算SSIM和Dice指标,代码默认从scripts执行"
    )
    parser.add_argument("--base-dir", type=str, default="..", help="项目根目录路径")
    parser.add_argument(
        "--subjects",
        type=str,
        nargs="*",
        help="要处理的受试者ID列表，不指定则处理所有受试者",
    )
    parser.add_argument(
        "--methods",
        type=str,
        nargs="*",
        choices=["processed"],
        help="要处理的方法类型列表，不指定则处理所有方法",
    )
    parser.add_argument(
        "--create-dirs", action="store_true", help="为指定的受试者创建标准目录结构"
    )
    parser.add_argument(
        "--no-reslice", action="store_false", dest="reslice", help="跳过重采样步骤"
    )
    parser.add_argument(
        "--no-ssim", action="store_false", dest="calc_ssim", help="跳过SSIM计算"
    )
    parser.add_argument(
        "--no-dice", action="store_false", dest="calc_dice", help="跳过Dice计算"
    )
    parser.add_argument(
        "--dry-run", action="store_true", help="模拟运行，不实际执行操作"
    )
    parser.add_argument("--verbose", "-v", action="store_true", help="详细输出")
    return parser.parse_args()


def main():
    """主函数"""
    args = parse_args()

    # 设置日志级别
    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)

    # 初始化处理器
    processor = MedicalImageProcessor(args.base_dir)

    # 如果只是创建目录结构
    if args.create_dirs:
        if not args.subjects:
            logging.error("创建目录时必须指定受试者ID列表 (--subjects)")
            return

        logging.info("执行目录创建操作...")
        if processor.create_subject_directories(args.subjects, args.dry_run):
            logging.info("目录创建操作完成!")
        else:
            logging.error("目录创建操作失败!")
        return

    # 确定要处理的受试者
    if args.subjects:
        subject_ids = args.subjects
    else:
        subject_ids = processor.get_subjects()

    if not subject_ids:
        logging.error("未找到受试者数据")
        return

    logging.info(f"将处理以下受试者: {', '.join(subject_ids)}")

    # 处理每个受试者
    success_count = 0
    for subject_id in subject_ids:
        if processor.process_subject(
            subject_id,
            args.reslice,
            args.calc_ssim,
            args.calc_dice,
            args.methods,
            args.dry_run,
        ):
            success_count += 1

    logging.info(f"处理完成! 成功处理 {success_count}/{len(subject_ids)} 个受试者")


if __name__ == "__main__":
    main()
